import sqlite3

class DatabaseManager:
    def __init__(self, db_name='yetian.db'):
        self.db_name = db_name
        self.conn = None
        self.cursor = None
        # 自动连接数据库
        self.connect()
    
    def connect(self):
        """连接到SQLite数据库"""
        try:
            self.conn = sqlite3.connect(self.db_name)
            self.cursor = self.conn.cursor()
            # 检查并添加media_code字段（如果不存在）
            self._ensure_media_code_field()
            return True
        except sqlite3.Error as e:
            print(f"数据库连接失败: {e}")
            return False
    
    def _ensure_media_code_field(self):
        """确保news表有media_code字段"""
        try:
            # 检查media_code字段是否存在
            self.cursor.execute("PRAGMA table_info(news)")
            columns = [column[1] for column in self.cursor.fetchall()]
            
            if 'media_code' not in columns:
                # 添加media_code字段
                self.execute("ALTER TABLE news ADD COLUMN media_code TEXT")
                print("已为news表添加media_code字段")
        except sqlite3.Error as e:
            print(f"检查/添加media_code字段时出错: {e}")
    
    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
            self.conn = None
            self.cursor = None
    
    def execute(self, sql, params=None):
        """执行SQL语句"""
        if not self.conn:
            if not self.connect():
                return False
        
        try:
            if params:
                self.cursor.execute(sql, params)
            else:
                self.cursor.execute(sql)
            self.conn.commit()
            return True
        except sqlite3.Error as e:
            print(f"SQL执行错误: {e}")
            self.conn.rollback()
            return False
    
    def fetchone(self):
        """获取单条结果"""
        try:
            return self.cursor.fetchone()
        except sqlite3.Error as e:
            print(f"获取结果失败: {e}")
            return None
    
    def fetchall(self):
        """获取所有结果"""
        try:
            return self.cursor.fetchall()
        except sqlite3.Error as e:
            print(f"获取结果失败: {e}")
            return []
            
    def insert_news(self, order_num, title, publish_at, url, app_id, face_img, media_code=None):
        """
        向news表插入一条新闻数据
        :param order_num: 新闻编号
        :param title: 标题
        :param publish_at: 发布时间
        :param url: 网址
        :param app_id: 应用ID
        :param face_img: 封面图片
        :param media_code: 媒体代码
        :return: (是否成功, 新闻ID)
        """
        try:
            # 插入数据，包含media_code字段
            sql = """
            INSERT INTO news (order_num, title, publish_at, url, app_id, face_img, media_code)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """
            params = (order_num, title, publish_at, url, app_id, face_img, media_code)
            
            if self.execute(sql, params):
                # 获取插入的ID
                news_id = self.cursor.lastrowid
                return True, news_id
            else:
                return False, None
        except Exception as e:
            print(f"插入新闻数据时出错: {e}")
            return False, None

# 创建news表（如果不存在）
def create_news_table():
    db = DatabaseManager('yetian.db')
    if db.connect():
        # 定义表结构
        create_table_sql = '''
        CREATE TABLE IF NOT EXISTS news (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER DEFAULT 0,
            order_num INTEGER DEFAULT 0,
            title TEXT NOT NULL,
            content TEXT,
            url TEXT,
            face_img TEXT,
            status INTEGER DEFAULT 1,
            publish_at TEXT DEFAULT CURRENT_TIMESTAMP,
            is_del INTEGER DEFAULT 0
        )
        '''
        
        if db.execute(create_table_sql):
            print("news表创建成功")
        else:
            print("news表创建失败")
        
        db.close()

# 插入一些测试数据（如果表为空）
def insert_test_data():
    db = DatabaseManager('yetian.db')
    if db.connect():
        # 检查是否已有数据
        db.execute("SELECT COUNT(*) FROM news")
        count = db.fetchone()[0]
        
        if count == 0:
            # 插入测试数据
            test_data = [
                (0, 1, "测试新闻标题1", "这是第一条测试新闻的内容", "https://example.com/news1", None, 1, "2023-10-01 10:00:00", 0),
                (0, 2, "测试新闻标题2", "这是第二条测试新闻的内容", "https://example.com/news2", None, 1, "2023-10-02 11:00:00", 0),
                (0, 3, "测试新闻标题3", "这是第三条测试新闻的内容", "https://example.com/news3", None, 1, "2023-10-03 12:00:00", 0)
            ]
            
            insert_sql = '''
            INSERT INTO news (user_id, order_num, title, content, url, face_img, status, publish_at, is_del)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            '''
            
            for data in test_data:
                if not db.execute(insert_sql, data):
                    print(f"插入测试数据失败: {data}")
                    break
            else:
                print("所有测试数据插入成功")
        else:
            print("news表已有数据，跳过插入测试数据")
        
        db.close()

# 如果直接运行此脚本，则创建表并插入测试数据
if __name__ == '__main__':
    create_news_table()
    insert_test_data()